import enum
import math
import time
from datetime import datetime
from typing import Callable, List

from rtde_control import RTDEControlInterface as RTDEControl
from rtde_io import RTDEIOInterface as RTDEIO
from rtde_receive import RTDEReceiveInterface as RTDEReceive

from utils.z_theta_and_Rxyz import Rxyz_to_theta, theta_to_Rxyz

from configs import cfgs


def log(msg):
    now = datetime.now()
    print(f"[{now:%H}:{now:%M}:{now:%S}] {msg}")


class Unit(enum.Enum):
    meter = 1
    millimeter = 1000


class Point:
    def __init__(self, x, y, z, unit:Unit=Unit.millimeter):
        self.x = x / unit.value
        self.y = y / unit.value
        self.z = z / unit.value

        self.unit = unit

    @property
    def xyz(self):
        return [self.x, self.y, self.z]

    def z_offset(self, distance=None):

        if distance is None:
            distance = 0.35
        else:
            distance = distance/self.unit.value

        return Point(self.x, self.y, self.z + distance, unit=Unit.meter)


class UnitVector(Point):
    def __init__(self, x, y, z, degree, unit=Unit.millimeter):
        Point.__init__(self, x, y, z, unit)
        self.degree = degree


def set_number(hint:str, default):
    while True:
        num = input(hint)
        if num == "":
            return default
        try:
            num = float(num)
            return num
        except:
            continue

class Task:

    v_fast = 0.5
    v_slow = 0.1
    v_very_slow = 0.02

    a_fast = 0.5
    a_slow = 0.2
    a_very_slow = 0.05

    blend_end = 0.0
    blend_mid = 0.02

    def __init__(self, robot_host):
        self.rtde_c = RTDEControl(robot_host)
        self.rtde_r = RTDEReceive(robot_host)
        self.rtde_io = RTDEIO(robot_host)
        self.rtde_receive = RTDEReceive(robot_host)

        self.z_angle_global = 0

        log("机器人已连接")
        log("")

    def make_path(self, point:Point, blend_mode, v=0.5, a=0.5):
        Rx, Ry, Rz = theta_to_Rxyz(self.z_angle_global / 180.0 * math.pi)
        return [point.x, point.y, point.z, Rx, Ry, Rz, v, a, blend_mode]

    def init_place(self):
        pass

    def pick_template(self):
        log("模拟电磁铁充电，吸起模板")
        time.sleep(1)

    def set_template(self, brick):
        log(f"模拟电磁铁放电，调整砖块{brick}的位置")
        time.sleep(1)

    def test(self):

        init = Point(x=0.000000, y=-0.636709, z=0.301307)
        piles = Point(x=0.698858, y=-0.030984, z=0.020000)
        bricks = [
            Point(x=-0.492987, y=-0.656833, z=0.020000),
            Point(x=-0.255655, y=-0.630635, z=0.020000),
            Point(x=-0.053405, y=-0.630635, z=0.020000),
            Point(x=0.116743, y=-0.623614, z=0.020000),
            Point(x=0.304833, y=-0.623614, z=0.020000),
            Point(x=0.492924, y=-0.623614, z=0.020000),
        ]

        # init place
        self.rtde_c.moveL(
            [self.make_path(init, self.blend_end)]
        )
        self.rtde_c.stopScript()
        log("机械臂在初始位置")

        time.sleep(1)

        for i, brick in enumerate(bricks):
            self.rtde_c.moveL([
                self.make_path(piles.z_offset(), self.blend_mid),
                self.make_path(piles, self.blend_end, v=self.v_slow, a=self.a_slow),
            ])

            self.pick_template()

            self.rtde_c.moveL([
                self.make_path(piles.z_offset(), self.blend_mid),
                self.make_path(brick.z_offset(), self.blend_mid),
                self.make_path(brick, self.blend_end, v=self.v_slow, a=self.a_slow),
            ])

            self.set_template(i+1)

            self.rtde_c.moveL([
                self.make_path(brick.z_offset(), self.blend_mid),
            ])

            log(f"砖块 {i+1} 放置程序完毕, 进行检查...")
            log("")
            time.sleep(1)

        self.rtde_c.stopScript()

        log("")
        log("测试程序完成")



    def use_magnet_tool(self, enable: bool, wait:bool=True):
        self.rtde_io.setStandardDigitalOut(6, enable)
        self.rtde_io.setStandardDigitalOut(7, enable)

        time.sleep(0.2)

        if self.rtde_receive.getDigitalOutState(6):
            log("Standard digital out (6) is HIGH")
        else:
            log("Standard digital out (6) is LOW")

        if self.rtde_receive.getDigitalOutState(7):
            log("Standard digital out (7) is HIGH")
        else:
            log("Standard digital out (7) is LOW")

        if wait:
            time.sleep(2)

    def test_magnet(self):
        while True:
            flag = input(" 1/0 > ")
            if flag == "1":
                log("magnet set to True")
                self.use_magnet_tool(True)
            elif flag == "0":
                log("magnet set to False")
                self.use_magnet_tool(False)


    def test_2(self):
        """
        测试控制工具头电磁铁，并且测试“试探性下降”
        :return:
        """

        self.use_magnet_tool(False, wait=False)

        const_accurate_z = -51.78
        const_target_z_offset = 0 # -30

        print(f"const_accurate_z: {const_accurate_z}")
        print(f"const_target_z_offset: {const_target_z_offset}")
        input("已就绪 > ")

        target = Point(152.77, -887.06, const_accurate_z)
        init_xyz = target.z_offset(300)

        self.rtde_c.moveL([
            self.make_path(init_xyz, self.blend_mid),
            self.make_path(target.z_offset(100), a=self.a_slow, blend_mode=self.blend_mid),
            self.make_path(target.z_offset(const_target_z_offset), a=self.a_very_slow, v=self.v_very_slow, blend_mode=self.blend_end),
        ])

        input("检查机械臂是否就位 > ")

        self.use_magnet_tool(True)

        self.rtde_c.moveL([
            self.make_path(target.z_offset(0.5), v=self.v_slow, a=self.a_slow, blend_mode=self.blend_end),
        ])

        input("执行放回动作 > ")

        self.rtde_c.moveL([
            self.make_path(target.z_offset(0.1), a=self.a_very_slow, blend_mode=self.blend_mid),
            self.make_path(target.z_offset(const_target_z_offset), a=self.a_very_slow, v=self.v_very_slow, blend_mode=self.blend_end),
        ])

        input("检查机械臂是否就位(放回) > ")
        self.use_magnet_tool(False)

        self.rtde_c.moveL([
            self.make_path(init_xyz, blend_mode=self.blend_end),
        ])

        log("")
        log("测试2结束")



    def test_final(self):

        const_accurate_z = -51.78
        const_target_z_offset = 0  # -30

        print(f"const_accurate_z: {const_accurate_z}")
        print(f"const_target_z_offset: {const_target_z_offset}")

        pile_center = Point(x=-114.278964, y=-788.496465, z=-51.780000)
        self.rtde_c.moveL([
            self.make_path(pile_center.z_offset(300), a=self.a_slow, v=self.v_slow, blend_mode=self.blend_end),
        ])

        input("已就绪 中心点校准程序还没有 > ")

        def down_and_up(base_point: Point, use_tool_callback: Callable, target_point: Point=None):

            if target_point is None:
                target_point = base_point

            self.rtde_c.moveL([
                self.make_path(base_point.z_offset(300), self.blend_end),
                self.make_path(base_point.z_offset(100), a=self.a_slow, blend_mode=self.blend_end),
                self.make_path(
                    base_point.z_offset(const_target_z_offset),
                    a=self.a_very_slow, v=self.v_very_slow, blend_mode=self.blend_end
                ),
            ])

            use_tool_callback()

            self.rtde_c.moveL([
                self.make_path(base_point.z_offset(100), v=self.v_slow, a=self.a_slow, blend_mode=self.blend_mid),
                self.make_path(target_point.z_offset(300), a=self.a_slow, blend_mode=self.blend_end),
            ])

        bricks: List[UnitVector] = [
            UnitVector(x=415.446518, y=-1014.388673, z=-51.780000, degree=15),
            UnitVector(x=178.114081, y=-1040.586114, z=-51.780000, degree=33),
            UnitVector(x=-24.135758, y=-1040.586114, z=-51.780000, degree=42),
            UnitVector(x=-194.283378, y=-1047.607612, z=-51.780000, degree=45),
            UnitVector(x=-382.373781, y=-1047.607612, z=-51.780000, degree=45),
            UnitVector(x=-570.464185, y=-1047.607612, z=-51.780000, degree=45),
        ]

        for i, brick in enumerate(bricks):

            down_and_up(pile_center, lambda: self.use_magnet_tool(True), brick)

            # rotate
            self.z_angle_global = brick.degree

            down_and_up(brick, lambda: self.use_magnet_tool(False), pile_center)

            # rotate back
            self.z_angle_global = 0

            log(f"砖块 {i+1} 放置程序完毕, 进行检查...")
            log("")
            time.sleep(1)


    def test_rotate(self):

        pile_center = Point(x=-114.278964, y=-788.496465, z=-51.780000)

        def act(base_point: Point):

            self.rtde_c.moveL([
                self.make_path(base_point.z_offset(0), blend_mode=self.blend_end),
            ])

        while True:
            # rotate
            self.z_angle_global = float(input("degree > "))
            act(pile_center)


    def test_robot_initialization(self):
        pass


    def test_move_until_contact(self):

        pile_center = Point(x=-114.278964, y=-788.496465, z=-51.780000)
        self.rtde_c.moveL([
            self.make_path(pile_center.xyz, )
        ])


if __name__ == '__main__':
    task = Task(robot_host=cfgs.robot_host)
    task.test_final()
